from torch import nn
import torch
from data import dataset
from visdom import Visdom

class Logistic(nn.Module):
    def __init__(self):
        super(Logistic, self).__init__()
        self.fc = nn.Linear(2, 1)
        self.sigmod = nn.Sigmoid()

    def forward(self, x):
        x = self.fc(x)
        x = self.sigmod(x)
        return x


lr = 0.001
optimizer_name = ["SGD", "Momentum", "RMSProp", "Adam"]
for name in optimizer_name:
    globals()[f'model_{name}'] = Logistic()
SGD = torch.optim.SGD(model_SGD.parameters(), lr=lr)
Momentum = torch.optim.SGD(model_Momentum.parameters(), lr=lr, momentum=0.8)
RMSProp = torch.optim.RMSprop(model_RMSProp.parameters(), lr=lr, alpha=0.9)
Adam = torch.optim.Adam(model_Adam.parameters(), lr=lr, betas=(0.9, 0.99))
criterion = nn.BCELoss()
data, labels = tuple(map(torch.tensor,  dataset.loadDataSet("./data/data.txt")))

viz = Visdom(env="optimizer")
loss_data = []
time_data = []
for j in range(10000):
    for name in optimizer_name:
        eval(f"{name}.zero_grad()")
    for i, name in enumerate(optimizer_name):
        globals()[f"out_{i}"] = eval(f'model_{name}')(data)
    temp = []
    for i, name in enumerate(optimizer_name):
        globals()[f"loss_{name}"] = criterion(globals()[f"out_{i}"], labels.unsqueeze(dim=1))
        temp.append(float(globals()[f"loss_{name}"].data))
    loss_data.append(temp)
    time_data.append([j for i in range(len(optimizer_name))])
    for name in optimizer_name:
        globals()[f"loss_{name}"].backward()
    for name in optimizer_name:
        globals()[name].step()
viz.line(Y=loss_data,
         X=time_data,
         win='line',
         opts=dict(legend=optimizer_name,
                   title='line test',
                   width=800,
                   height=800,
                   xlabel='Time',
                   ylabel='Volume'),
         update='append'
         )






